package rtpps

import (
	"sync"
	"time"

	"github.com/pion/rtp"
	"github.com/yapingcat/gomedia/mpeg2"
	"gogs.smarteye.com/smarteye/gortc/pkg/format/rtpmpeg4video"
)

type Encoder struct {
	*rtpmpeg4video.Encoder

	mutex    sync.Mutex
	muxer    *mpeg2.PSMuxer
	cidToPid map[mpeg2.PS_STREAM_TYPE]uint8
}

// Init initializes the encoder.
func (e *Encoder) Init() error {
	if err := e.Encoder.Init(); err != nil {
		return err
	}

	e.muxer = mpeg2.NewPsMuxer()
	e.cidToPid = make(map[mpeg2.PS_STREAM_TYPE]uint8)

	return nil
}

// Encode encodes a frame into RTP packets.
func (e *Encoder) Encode(frame []byte, cid mpeg2.PS_STREAM_TYPE, pts time.Duration) ([]*rtp.Packet, error) {
	e.mutex.Lock()
	defer e.mutex.Unlock()

	pid, ok := e.cidToPid[cid]
	if !ok {
		pid = e.muxer.AddStream(cid)
		e.cidToPid[cid] = pid
	}

	var result []*rtp.Packet
	e.muxer.OnPacket = func(pkg []byte) {
		if packets, _ := e.Encoder.Encode(pkg); len(packets) > 0 {
			result = append(result, packets...)
		}
	}

	iPTS := uint64(pts.Microseconds())
	if err := e.muxer.Write(pid, frame, iPTS, iPTS); err != nil {
		return nil, err
	}

	return result, nil
}
